# =============================================================================
# PROJECT CHRONO - http://projectchrono.org
#
# Copyright (c) 2014 projectchrono.org
# All rights reserved.
#
# Use of this source code is governed by a BSD-style license that can be found
# in the LICENSE file at the top level of the distribution and at
# http://projectchrono.org/license-chrono.txt.
#
# =============================================================================

import pychrono.core as chrono
import pychrono.irrlicht as chronoirr
import pychrono.cascade as cascade
from OCC.Core import TopoDS

print ("Example: Load a STEP file, generated by a CAD")

#  Create the simulation system and add items
sys = chrono.ChSystemNSC()

# Load a STEP file, containing a mechanism. 
# The demo STEP file has been created using a 3D CAD (in this case, SolidEdge v.18).

# Create the ChCascadeDoc, a container that loads the STEP model and manages its subassembles
mydoc = cascade.ChCascadeDoc()

# Load the STEP model
load_ok = mydoc.LoadSTEP(chrono.GetChronoDataFile('cascade/assembly.stp'))  

# In most CAD systems the Y axis is horizontal, but we want it vertical.
# So define a root transformation for rotating all the imported objects.
rotation1 = chrono.ChQuaterniond()
rotation1.SetFromAngleAxis(-chrono.CH_PI_2, chrono.ChVector3d(1, 0, 0))  # 1: rotate 90 deg on X axis
rotation2 = chrono.ChQuaterniond()
rotation2.SetFromAngleAxis(chrono.CH_PI, chrono.ChVector3d(0, 1, 0))  # 2: rotate 180 deg on vertical Y axis
tot_rotation = chrono.ChQuaterniond()
tot_rotation = rotation2 * rotation1     # rotate on 1 then on 2, using quaternion product
root_frame = chrono.ChFrameMovingd(chrono.ChVector3d(0, 0, 0), tot_rotation)

# Retrieve some sub shapes from the loaded model, using
# the GetNamedShape() function, that can use path/subpath/subsubpath/part
# syntax and * or ? wildcards, etc.
rigidBody1 = 0
rigidBody2 = 0

if load_ok:     
    shape1 = TopoDS.TopoDS_Shape()
    if (mydoc.GetNamedShape(shape1, "Assem1/body1")):
        body1 = cascade.ChCascadeBodyEasy(shape1, 1000) # density
        body1.SetFixed(True) 
        sys.Add(body1)
        
        # Move the body as for global displacement/rotation (also body1 %= root_frame )
        body1.ConcatenatePreTransformation(root_frame)
        
        rigidBody1 = body1
    else:
        print("Warning. Desired object not found in document\n")

    shape2 = TopoDS.TopoDS_Shape()
    if (mydoc.GetNamedShape(shape2, "Assem1/body2")): 
        body2 = cascade.ChCascadeBodyEasy(shape2, 1000) # density        
        sys.Add(body2)
        
        # Move the body as for global displacement/rotation  (also mbody2 %= root_frame )
        body2.ConcatenatePreTransformation(root_frame)

        rigidBody2 = body2
    else:
        print("Warning. Desired object not found in document\n")

else:
    print("Warning. Desired STEP file could not be opened/parsed \n")

# Create a revolute joint between the two parts
# as in a pendulum. We assume we already know in advance
# the aboslute position of the joint (ex. we used measuring tools in the 3D CAD)

measured_joint_pos_mm = chrono.ChVector3d(0, 48, 120)

scale = 1. / 1000.  # because we use meters instead of mm

joint_pos = chrono.ChVector3d(root_frame.TransformPointLocalToParent(measured_joint_pos_mm * scale)) 

if (rigidBody1 and rigidBody2):
    link = chrono.ChLinkLockRevolute()
    link.Initialize(rigidBody1, rigidBody2, chrono.ChFramed(joint_pos))
    sys.Add(link)

# Create a large cube as a floor.
floor = chrono.ChBodyEasyBox(1, 0.2, 1, 1000)
floor.SetPos(chrono.ChVector3d(0,-0.3,0))
floor.SetFixed(True)
floor.GetVisualShape(0).SetTexture(chrono.GetChronoDataFile('textures/blue.png'))
sys.Add(floor)

# Create an Irrlicht application to visualize the system
vis = chronoirr.ChVisualSystemIrrlicht()
vis.AttachSystem(sys)
vis.SetWindowSize(1024,768)
vis.SetWindowTitle('Test')
vis.Initialize()
vis.AddLogo(chrono.GetChronoDataFile('logo_chrono_alpha.png'))
vis.AddSkyBox()
vis.AddCamera(chrono.ChVector3d(0.2,0.2,-0.2))
vis.AddTypicalLights()

# Run the simulation
timestep = 0.01
rt_timer = chrono.ChRealtimeStepTimer()

while vis.Run():
    vis.BeginScene()
    vis.Render()
    vis.EndScene()
    sys.DoStepDynamics(timestep)
    rt_timer.Spin(timestep)